package com.ncepu.jjojcodesandbox.interceptor;

import com.ncepu.jjojcodesandbox.exception.RateLimitException;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.stereotype.Component;
import org.springframework.web.servlet.HandlerInterceptor;

import javax.annotation.Resource;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.util.concurrent.TimeUnit;

@Component
public class RateLimitInterceptor implements HandlerInterceptor {

    @Resource
    private StringRedisTemplate stringRedisTemplate;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) {
        // 获取客户端 IP 地址
        String ip = getClientIp(request);
        String rateKey = "rate_limit:ip:" + ip;

        Long count = stringRedisTemplate.opsForValue().increment(rateKey);
        if (count == null) {
            return true;
        }
        if (count == 1) {
            stringRedisTemplate.expire(rateKey, 1, TimeUnit.MINUTES);
        }
        if (count > 5) {
            response.setStatus(429);
            throw new RateLimitException("请求过于频繁");
        }
        return true;
    }

    /**
     * 获取客户端 IP 地址
     * @param request
     * @return
     */
    public String getClientIp(HttpServletRequest request) {
        String ip = request.getHeader("X-Forwarded-For");
        if (ip != null && ip.length() != 0 && !"unknown".equalsIgnoreCase(ip)) {
            // 如果有多个 IP，取第一个
            return ip.split(",")[0];
        }
        ip = request.getHeader("Proxy-Client-IP");
        if (ip != null && ip.length() != 0 && !"unknown".equalsIgnoreCase(ip)) {
            return ip;
        }
        ip = request.getHeader("WL-Proxy-Client-IP");
        if (ip != null && ip.length() != 0 && !"unknown".equalsIgnoreCase(ip)) {
            return ip;
        }
        ip = request.getHeader("HTTP_CLIENT_IP");
        if (ip != null && ip.length() != 0 && !"unknown".equalsIgnoreCase(ip)) {
            return ip;
        }
        ip = request.getHeader("HTTP_X_FORWARDED_FOR");
        if (ip != null && ip.length() != 0 && !"unknown".equalsIgnoreCase(ip)) {
            return ip;
        }
        // 如果没有通过代理，则直接获取远程地址
        return request.getRemoteAddr();
    }
}
